import math
import random
from collections import defaultdict


def euclidean_distance(point1, point2):
    return math.sqrt(sum((a - b) ** 2 for a, b in zip(point1, point2)))


def assign_cluster(x, centroids):
    distances = [euclidean_distance(x, centroid) for centroid in centroids]
    return distances.index(min(distances))


def initialize_centroids(data, k):
    return random.sample(data, k)


def update_centroids(data, cluster_assignments, k):
    clusters = defaultdict(list)
    for i, cluster_idx in enumerate(cluster_assignments):
        clusters[cluster_idx].append(data[i])

    new_centroids = []
    for i in range(k):
        if i in clusters and clusters[i]:
            cluster_points = clusters[i]
            dimension = len(cluster_points[0])
            new_centroid = [
                sum(point[d] for point in cluster_points) / len(cluster_points)
                for d in range(dimension)
            ]
            new_centroids.append(new_centroid)
        else:
            new_centroids.append(random.choice(data))

    return new_centroids


def has_converged(old_centroids, new_centroids, epsilon):
    for old, new in zip(old_centroids, new_centroids):
        if euclidean_distance(old, new) > epsilon:
            return False
    return True


def Kmeans(data, k, epsilon=1e-4, iteration=100):
    if len(data) < k:
        raise ValueError("数据点数量不能小于聚类数量k")

    if k <= 0:
        raise ValueError("聚类数量k必须大于0")

    centroids = initialize_centroids(data, k)

    for iter_count in range(iteration):
        cluster_assignments = [assign_cluster(point, centroids) for point in data]

        new_centroids = update_centroids(data, cluster_assignments, k)

        if has_converged(centroids, new_centroids, epsilon):
            print(f"算法在第 {iter_count + 1} 次迭代后收敛")
            break

        centroids = new_centroids

    else:
        print(f"达到最大迭代次数 {iteration}")

    return centroids, cluster_assignments

if __name__ == "__main__":
    random.seed(42)  # 设置随机种子以便复现结果

    data = []
    for _ in range(30):
        data.append([random.gauss(2, 0.5), random.gauss(2, 0.5)])
    for _ in range(30):
        data.append([random.gauss(8, 0.5), random.gauss(8, 0.5)])
    for _ in range(30):
        data.append([random.gauss(5, 0.5), random.gauss(2, 0.5)])

    print("数据集大小:", len(data))
    print("前5个数据点:", data[:5])

    centroids, assignments = Kmeans(data, k=3, epsilon=0.001, iteration=100)

    print("\n聚类结果:")
    print("聚类中心:")
    for i, centroid in enumerate(centroids):
        print(f"簇{i}: {centroid}")

    cluster_sizes = defaultdict(int)
    for assignment in assignments:
        cluster_sizes[assignment] += 1

    print("\n各簇大小:")
    for cluster_idx, size in sorted(cluster_sizes.items()):
        print(f"簇{cluster_idx}: {size}个点")

    test_point = [3, 3]
    nearest_cluster = assign_cluster(test_point, centroids)
    print(f"\n测试点 {test_point} 被分配到簇 {nearest_cluster}")